# 最短路径Dijkstra算法实现代码

# 定义顶点的访问状态：
# RED(红)：已访问且已确定最短路径
# GREEN(绿): 已访问但未确定当前长度为最短
# WHITE(白): 未访问
RED = 2
GREEN = 1
WHITE = 0

## 计算最短路径
## n 表示最多总共有n个顶点
## edges 所有路径信息，以二维数组给出，
##     数组元素edges[i] = [s, t, w] 表示从 s 到 t 存在权值为 w 的路径
## src 起点编号
## 返回：输出从起点到所有顶点的路径长度，并返回最后一个路径长度
def findAllShortestPath(n: int, edges: list[list[int]], src: int) -> int:
    # dist[x]: 保存起点到 x 的最短路径长度
    # 初始值为理论上的最大值（所有边长加起来再加1）
    INF = sum(e[2] for e in edges) + 1
    dist = [INF] * (n + 1) 

    # visit[x]表示 x 的访问状态：
    #     0(白)-未访问(初始值) 
    #     1(绿)-已访问未确定 
    #     2(红)-已确定最短长度
    visit = [WHITE] * (n + 1)

    # 初始时只有起点被访问到，设置路径长度为0
    visit[src] = visit[0] = GREEN
    dist[src] = 0

    # 执行dijkstra搜索最短路径
    for _ in range(n):
        # 搜索所有的绿点，找到路径最短的一个
        # 标记为红点，并作为当前搜索点
        s = min([x for x in range(0, n + 1) if visit[x] == GREEN], key=lambda x: dist[x])

        # 如果没有更多可访问的顶点则直接退出
        if (s == 0): break
        
        # 标记 s 为当前顶点，设置访问状态为红色并开始搜索
        visit[s] = RED

        # 下面这行代码输出中间过程，可用于调试
        print('s=', s, dist[s])

        # 枚举当前顶点 s 的邻接路径，发现新的路径并更新维护排序数组
        for u, v, w in edges:
            # 当前边是 s 的邻边，且目标顶点未完成搜索（不是红色）
            if s == u and visit[v] != RED:
                # 计算从起点经过s到达v的路径长度
                d = dist[s] + w
                # 计算出一条更短的路径长度，更新路径长度及访问状态
                if d < dist[v]:
                    dist[v] = d
                    visit[v] = GREEN

    # 输出结果并返回
    print('所有的顶点路径信息:')
    for i, d in enumerate(dist):
        if visit[i] != RED: continue
        print('目标顶点', i, '->', d)
    return max(dist[1:])

# main
n, src = 5, 1
edges = [[1,2,2],[1,5,8],[1,4,3],[2,3,5],[3,5,1],[4,3,2],[4,5,6]]
ans = findAllShortestPath(n, edges, src)
print('搜索完毕，最远顶点路径为:', ans)
